open Spitree;;

(* SMap is the basis for the substitution map of the
 * alpha conversion function. The substitution map
 * indicates what each name in the current scope must
 * be substituted with.
 *)
module SMap = Map.Make(String)

(* --- subst_name ---
 * Called on free names
 * Given some name, we do the following:
 * If the name is in the map of substitutions
 * (It is either bound, or an already found free variable):
 * - Substitute the name
 * Else (we have found a new free variable!):
 * - Create a new name, and update the map and accumulator
 *)
let subst_name x accum substMap =
  if SMap.mem x substMap then
    (SMap.find x substMap, accum, substMap)
  else
    let newname = (string_of_int accum) ^ "_" ^ x in
    let newMap = SMap.add x newname substMap in
    (newname, accum + 1, newMap)

(* --- bind ---
 * Called for the binding occurrence of a name:
 * We generate a new unique name and update
 * the substitution map.
 * We return the new name, an updated accumulator
 * and the updated map.
 * We also return a tuple (x, oldname) or (x, ""),
 * containing the mapping we replace in this scope,
 * if any.
 *)
let bind x accum substMap =
  let newname = (string_of_int accum) ^ "_" ^ x in
  let newMap = SMap.add x newname substMap in
  if SMap.mem x substMap then
    let oldname = SMap.find x substMap in
    (newname, accum + 1, newMap, (x, oldname))
  else
    (newname, accum + 1, newMap, (x, ""))

(* --- unbind ---
 * When exiting a scope, we must call unbind
 * with any mappings from the outer scope that must
 * be reinserted in the substitution map.
 *)
let unbind (x, oldname) substMap =
  if oldname = "" then
    SMap.remove x substMap
  else
    SMap.add x oldname substMap

(* Input: a - accumulator, s - substitution map *)
(* Returns: *_dat - updated datastructure, ua - updated accumulator, us - updated substitution map *)
let rec msg_subst msg a s = match msg with
  | Name (x) ->
      let (x_dat, ua, us) = (subst_name x a s) in
      (Name (x_dat), ua, us)
  | Encr (m, n) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (n_dat, ua, us) = (msg_subst n ua us) in
      (Encr (m_dat, n_dat), ua, us)
  | MPair (m, n) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (n_dat, ua, us) = (msg_subst n ua us) in
      (MPair (m_dat, n_dat), ua, us)
  | Ok -> (Ok, a, s);;

let rec reso_subst reso a s = match reso with
  | Reso (l, m) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      (Reso (l, m_dat), ua, us)
  | RPair (r, t) ->
      let (r_dat, ua, us) = (reso_subst r a s) in
      let (t_dat, ua, us) = (reso_subst t ua us) in
      (RPair (r_dat, t_dat), ua, us)
  | Empty -> (Empty, a, s);;

let rec typ_subst typ a s = match typ with
  | Un -> (Un, a, s)
  | Key (t) ->
      let (t_dat, ua, us) = (typ_subst t a s) in
      (Key (t_dat), ua, us)
  | Ch (t) ->
      let (t_dat, ua, us) = (typ_subst t a s) in
      (Ch (t_dat), ua, us)
  | TPair (x, t, u) ->
      let (t_dat, ua, us) = (typ_subst t a s) in
      let (x_dat, ua, us, binding) = (bind x ua us) in
      let (u_dat, ua, us) = (typ_subst u ua us) in
      (TPair (x_dat, t_dat, u_dat), ua, (unbind binding us))
  | TOk (r) ->
      let (r_dat, ua, us) = (reso_subst r a s) in
      (TOk (r_dat), ua, us)
  | Unknown -> (Unknown, a, s);;

let rec proc_subst proc a s = match proc with
  | Out (m, n) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (n_dat, ua, us) = (msg_subst n ua us) in
      (Out (m_dat, n_dat), ua, us)
  | In (m, x, p) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (x_dat, ua, us, binding) = (bind x ua us) in
      let (p_dat, ua, us) = (proc_subst p ua us) in
      (In (m_dat, x_dat, p_dat), ua, (unbind binding us))
  | InRepl (m, x, p) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (x_dat, ua, us) = (subst_name x ua us) in
      let (p_dat, ua, us) = (proc_subst p ua us) in
      (InRepl (m_dat, x_dat, p_dat), ua, us)
  | Nu (x, t, p) ->
      let (t_dat, ua, us) = (typ_subst t a s) in
      let (x_dat, ua, us, binding) = (bind x ua us) in
      let (p_dat, ua, us) = (proc_subst p ua us) in
      (Nu (x_dat, t_dat, p_dat), ua, (unbind binding us))
  | PPara (p, q) ->
      let (p_dat, ua, us) = (proc_subst p a s) in
      let (q_dat, ua, us) = (proc_subst q ua us) in
      (PPara (p_dat, q_dat), ua, us)
  | PNil -> (PNil, a, s)
  | Begin (l, m) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      (Begin (l, m_dat), ua, us)
  | End (l, m) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      (End (l, m_dat), ua, us)
  | Decrypt (m, x, t, n, p) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (t_dat, ua, us) = (typ_subst t ua us) in
      let (n_dat, ua, us) = (msg_subst n ua us) in
      let (x_dat, ua, us, binding) = (bind x ua us) in
      let (p_dat, ua, us) = (proc_subst p ua us) in
      (Decrypt (m_dat, x_dat, t_dat, n_dat, p_dat), ua, (unbind binding us))
  | Split (m, x, t, y, u, p) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (t_dat, ua, us) = (typ_subst t ua us) in
      let (u_dat, ua, us) = (typ_subst u ua us) in
      let (x_dat, ua, us, xbinding) = (bind x ua us) in
      let (y_dat, ua, us, ybinding) = (bind y ua us) in
      let (p_dat, ua, us) = (proc_subst p ua us) in
      (Split (m_dat, x_dat, t_dat, y_dat, u_dat, p_dat), ua, (unbind xbinding (unbind ybinding us)))
  | Match (m, t, n, x, usub, p) ->
      let (m_dat, ua, us) = (msg_subst m a s) in
      let (t_dat, ua, us) = (typ_subst t ua us) in
      let (n_dat, ua, us) = (msg_subst n ua us) in
      let (usub_dat, ua, us) = (typ_subst usub ua us) in
      let (x_dat, ua, us, binding) = (bind x ua us) in
      let (p_dat, ua, us) = (proc_subst p ua us) in
      (Match (m_dat, t_dat, n_dat, x_dat, usub_dat, p_dat), ua, (unbind binding us));;
